import requests
import pandas as pd
import io
import os
from datetime import datetime
from loguru import logger

class LightCurveDownloader:
    def __init__(self, base_url, save_dir="downloaded_light_curves", query_defaults=None, timeout=30):
        """
        初始化Light Curve Downloader

        :param base_url: API的基础URL
        :param save_dir: 保存文件夹
        :param query_defaults: 查询默认参数(dict)
        :param timeout: 请求超时时间(秒)
        """
        self.base_url = base_url
        self.save_dir = save_dir
        self.query_defaults = query_defaults or {}
        self.timeout = timeout

        os.makedirs(self.save_dir, exist_ok=True)
        logger.info(f"保存目录设置为：{self.save_dir}")

    def query_single_target(self, ra, dec):
        """
        查询单个目标光变曲线
        
        :param ra: 赤经，字符串格式
        :param dec: 赤纬，字符串格式
        :return: pandas.DataFrame 或 None
        """
        params = self.query_defaults.copy()
        params.update({"ra": ra, "dec": dec})
        
        try:
            response = requests.get(self.base_url, params=params, timeout=self.timeout)
            response.raise_for_status()

            logger.info(f"查询URL: {response.url}")

            content_type = response.headers.get('Content-Type', '')

            if 'application/json' in content_type:
                data = response.json()
                if not data:
                    logger.warning(f"[{ra}, {dec}] 返回为空")
                    return None
                df = pd.DataFrame(data)

            elif 'text/csv' in content_type or 'text/plain' in content_type:
                df = pd.read_csv(io.StringIO(response.text))

            else:
                logger.error(f"[{ra}, {dec}] 未知返回格式: {content_type}")
                return None

            return df

        except Exception as e:
            logger.error(f"[{ra}, {dec}] 查询失败: {e}")
            return None

    def sanitize_filename(self, ra, dec):
        """
        生成合法文件名
        """
        ra_clean = ra.replace(':', '-')
        dec_clean = dec.replace(':', '-').replace('+', 'p').replace('-', 'm')
        return f"lc_RA{ra_clean}_DEC{dec_clean}.csv"

    def query_all_targets(self, targets):
        """
        批量查询一组目标
        :param targets: 列表，每个元素是 {"ra": "XX", "dec": "XX"}
        """
        logger.info(f"开始批量查询，共 {len(targets)} 个目标")

        for idx, target in enumerate(targets, start=1):
            ra = target["ra"]
            dec = target["dec"]
            logger.info(f"({idx}/{len(targets)}) 查询目标 RA={ra} DEC={dec}")

            df = self.query_single_target(ra, dec)
            if df is not None and not df.empty:
                filename = self.sanitize_filename(ra, dec)
                save_path = os.path.join(self.save_dir, filename)
                df.to_csv(save_path, index=False)
                logger.success(f"已保存: {save_path}")
            else:
                logger.warning(f"目标 RA={ra} DEC={dec} 查询无数据")

        logger.info("全部目标查询完成！")
